#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
智能化漏洞检测与攻击路径分析平台 - 检查与运行脚本

此脚本用于检查系统依赖项和运行状态，确保项目能够正常运行。
它会验证关键组件的可用性，并测试健康检查端点。
"""

import os
import sys
import time
import subprocess
import json
import requests
from datetime import datetime

# 项目根目录
PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))

# 日志文件路径
LOG_FILE = os.path.join(PROJECT_ROOT, "check_log.txt")

# 服务端点
HEALTH_CHECK_ENDPOINT = "http://localhost:5000/health"
FLASK_ENDPOINT = "http://localhost:5000"
STREAMLIT_ENDPOINT = "http://localhost:8501"

# 服务端口
FLASK_PORT = 5000
STREAMLIT_PORT = 8501
REDIS_PORT = 6379

# 检查结果状态码
STATUS_OK = 0
STATUS_WARNING = 1
STATUS_ERROR = 2


def log(message, level="INFO"):
    """记录日志消息到控制台和文件"""
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    log_message = f"[{timestamp}] [{level}] {message}"
    
    # 输出到控制台
    if level == "ERROR":
        print("\033[91m" + log_message + "\033[0m")  # 红色
    elif level == "WARNING":
        print("\033[93m" + log_message + "\033[0m")  # 黄色
    else:
        print(log_message)  # 白色
    
    # 写入日志文件
    with open(LOG_FILE, "a", encoding="utf-8") as f:
        f.write(log_message + "\n")


def check_python_version():
    """检查Python版本"""
    log("检查Python版本...")
    required_min_version = (3, 8)
    current_version = sys.version_info[:2]
    
    if current_version < required_min_version:
        log(f"警告: Python版本 {'.'.join(map(str, current_version))} 低于推荐版本 {'.'.join(map(str, required_min_version))}", "WARNING")
        return STATUS_WARNING
    else:
        log(f"Python版本 {'.'.join(map(str, current_version))} 符合要求")
        return STATUS_OK


def check_redis():
    """检查Redis服务是否可用"""
    log("检查Redis服务...")
    try:
        # 尝试使用redis-cli检查
        subprocess.run(["redis-cli", "ping"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        log("Redis服务可用")
        return STATUS_OK
    except (subprocess.SubprocessError, FileNotFoundError):
        log("警告: Redis服务当前不可用或未安装。start.py脚本可能会尝试自动处理此问题。", "WARNING")
        log("提示: 对于Windows系统，您可以使用WSL安装Redis或使用Redis Windows版本。", "INFO")
        return STATUS_WARNING


def check_dependencies():
    """检查Python依赖项"""
    log("检查Python依赖项...")
    requirements_file = os.path.join(PROJECT_ROOT, "requirements.txt")
    
    if not os.path.exists(requirements_file):
        log(f"错误: 找不到requirements.txt文件在 {requirements_file}", "ERROR")
        return STATUS_ERROR
    
    try:
        with open(requirements_file, "r", encoding="utf-8") as f:
            dependencies = [line.strip() for line in f if line.strip() and not line.startswith("#")]
        
        missing_deps = []
        for dep in dependencies:
            # 尝试导入包（简化版本）
            # 这里只检查主要的包名，忽略版本号和额外要求
            try:
                package_name = dep.split("==")[0].split(">=")[0].split("<=")[0].strip()
                # 对于一些特殊情况进行处理
                if package_name in ["python-dotenv", "flask-cors"]:
                    package_name = package_name.replace("-", "_")
                __import__(package_name)
            except ImportError:
                missing_deps.append(dep)
        
        if missing_deps:
            log(f"警告: 缺少以下依赖项: {', '.join(missing_deps)}", "WARNING")
            log("建议运行: pip install -r requirements.txt", "WARNING")
            return STATUS_WARNING
        else:
            log("所有依赖项检查通过")
            return STATUS_OK
    except Exception as e:
        log(f"错误: 检查依赖项时发生错误: {str(e)}", "ERROR")
        return STATUS_ERROR


def check_ports():
    """检查必要的端口是否可用"""
    log("检查必要端口是否可用...")
    try:
        import socket
        
        ports_to_check = [("Flask API", FLASK_PORT), ("Streamlit", STREAMLIT_PORT)]
        occupied_ports = []
        
        for name, port in ports_to_check:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                result = s.connect_ex(("127.0.0.1", port))
                if result == 0:
                    occupied_ports.append((name, port))
        
        if occupied_ports:
            log("警告: 以下端口已被占用:", "WARNING")
            for name, port in occupied_ports:
                log(f"  - {name}: {port}", "WARNING")
            log("您可能需要关闭占用这些端口的服务或修改配置文件中的端口设置。", "WARNING")
            return STATUS_WARNING
        else:
            log("所有必要端口都可用")
            return STATUS_OK
    except Exception as e:
        log(f"错误: 检查端口时发生错误: {str(e)}", "ERROR")
        return STATUS_ERROR


def check_health_endpoint():
    """检查健康检查端点是否正常响应"""
    log("检查健康检查端点...")
    try:
        response = requests.get(HEALTH_CHECK_ENDPOINT, timeout=5)
        if response.status_code == 200:
            health_data = response.json()
            log(f"健康检查端点响应正常: 状态 - {health_data.get('status')}")
            
            # 检查关键组件状态
            system_status = health_data.get('system_status', {})
            for component, status in system_status.items():
                if status.lower() != "ok":
                    log(f"警告: {component} 状态异常: {status}", "WARNING")
                    return STATUS_WARNING
            
            return STATUS_OK
        else:
            log(f"错误: 健康检查端点返回非200状态码: {response.status_code}", "ERROR")
            return STATUS_ERROR
    except requests.exceptions.ConnectionError:
        log(f"错误: 无法连接到健康检查端点 {HEALTH_CHECK_ENDPOINT}。确保服务已启动。", "ERROR")
        return STATUS_ERROR
    except Exception as e:
        log(f"错误: 检查健康端点时发生错误: {str(e)}", "ERROR")
        return STATUS_ERROR


def run_start_script():
    """运行启动脚本"""
    log("准备运行启动脚本...")
    start_script = os.path.join(PROJECT_ROOT, "start.py")
    
    if not os.path.exists(start_script):
        log(f"错误: 找不到启动脚本 {start_script}", "ERROR")
        return STATUS_ERROR
    
    log(f"正在运行启动脚本: {start_script}")
    log("注意: 这将启动所有必要的服务。按Ctrl+C停止。")
    
    try:
        # 在新进程中运行启动脚本，这样它可以继续运行
        process = subprocess.Popen([sys.executable, start_script], 
                                  cwd=PROJECT_ROOT, 
                                  stdout=subprocess.PIPE, 
                                  stderr=subprocess.PIPE, 
                                  text=True)
        
        # 实时输出日志
        for line in process.stdout:
            print(line.strip())
        
        # 检查进程状态
        process.wait()
        
        if process.returncode != 0:
            log(f"错误: 启动脚本返回非零退出码: {process.returncode}", "ERROR")
            # 显示错误输出
            for line in process.stderr:
                print("\033[91m" + line.strip() + "\033[0m")
            return STATUS_ERROR
        else:
            log("启动脚本执行成功")
            return STATUS_OK
    except Exception as e:
        log(f"错误: 运行启动脚本时发生错误: {str(e)}", "ERROR")
        return STATUS_ERROR


def test_api_endpoints():
    """测试关键API端点"""
    log("测试关键API端点...")
    endpoints_to_test = [
        ("Flask根路径", FLASK_ENDPOINT),
        ("健康检查", HEALTH_CHECK_ENDPOINT)
    ]
    
    results = {}
    for name, endpoint in endpoints_to_test:
        try:
            response = requests.get(endpoint, timeout=5)
            results[name] = {
                "status": "OK" if response.status_code == 200 else f"FAIL (HTTP {response.status_code})",
                "code": response.status_code
            }
        except Exception as e:
            results[name] = {
                "status": f"ERROR: {str(e)}",
                "code": -1
            }
    
    # 显示测试结果
    all_ok = True
    for name, result in results.items():
        if "FAIL" in result["status"] or "ERROR" in result["status"]:
            log(f"{name}: {result['status']}", "ERROR")
            all_ok = False
        else:
            log(f"{name}: {result['status']}")
    
    return STATUS_OK if all_ok else STATUS_ERROR


def show_usage():
    """显示脚本使用说明"""
    print("\n智能化漏洞检测与攻击路径分析平台 - 检查与运行脚本")
    print("用法: python check_and_run.py [选项]")
    print("\n选项:")
    print("  --check-only     仅执行系统检查，不运行启动脚本")
    print("  --run-only       直接运行启动脚本，跳过检查")
    print("  --test-endpoints 测试API端点（要求服务已启动）")
    print("  --force          强制运行启动脚本，即使检查有警告")
    print("  --help           显示此帮助信息")
    print("\n无参数时：执行所有检查，然后运行启动脚本")


def main():
    """主函数"""
    # 清空旧日志
    open(LOG_FILE, "w").close()
    log("开始系统检查与运行脚本")
    
    # 解析命令行参数
    args = sys.argv[1:]
    check_only = "--check-only" in args
    run_only = "--run-only" in args
    test_endpoints = "--test-endpoints" in args
    force = "--force" in args
    show_help = "--help" in args
    
    if show_help:
        show_usage()
        return STATUS_OK
    
    # 运行测试端点模式
    if test_endpoints:
        return test_api_endpoints()
    
    # 运行检查模式
    if not run_only:
        log("\n========== 开始系统检查 ==========")
        checks = [
            ("Python版本", check_python_version),
            ("Redis服务", check_redis),
            ("Python依赖项", check_dependencies),
            ("端口可用性", check_ports)
        ]
        
        overall_status = STATUS_OK
        for name, check_func in checks:
            log(f"\n[{name}]")
            status = check_func()
            if status > overall_status:
                overall_status = status
        
        log("\n========== 系统检查完成 ==========")
        if overall_status == STATUS_OK:
            log("所有检查通过！系统准备就绪。")
        elif overall_status == STATUS_WARNING:
            log("检查完成，但存在警告。start.py脚本可能会尝试自动解决这些问题。", "WARNING")
            if not force and not check_only:
                log("您可以使用 --force 参数强制运行启动脚本。", "INFO")
        else:
            log("检查失败，存在错误。请解决这些错误后再运行系统。", "ERROR")
            log(f"详细日志请查看: {LOG_FILE}", "WARNING")
            if not force:
                return overall_status
    
    # 运行启动脚本
    if not check_only:
        log("\n========== 开始运行系统 ==========")
        log("注意: start.py脚本将尝试自动安装缺少的依赖并处理Redis服务")
        log("如果遇到问题，请参考详细日志和README文档")
        return run_start_script()
    
    return STATUS_OK


if __name__ == "__main__":
    exit_code = main()
    sys.exit(exit_code)